package mac5789.projeto4.model;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;

/**
 * Modelo do mundo declarável através de métodos de modelagem
 * 
 * @author lundberg
 * 
 */
public class DeclarableWorldModel implements WorldModel {

	private final Map<Object, State> states = new LinkedHashMap<Object, State>();

	private final Map<Object, Action> actions = new LinkedHashMap<Object, Action>();

	private final Map<TransitionModelKey, TransitionModel> transitions = new HashMap<DeclarableWorldModel.TransitionModelKey, TransitionModel>();

	private final Map<State, Set<Action>> availableActionsAt = new HashMap<State, Set<Action>>();

	private final DeclarableRewardModel rewardModel = new DeclarableRewardModel();

	private State initialState;

	private final Set<State> goals = new HashSet<State>();

	private int maxIteractions = Integer.MAX_VALUE;

	/**
	 * {@inheritDoc}
	 */
	@Override
	public Set<State> states() {
		return new HashSet<State>(states.values());
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public Set<Action> actions() {
		return new HashSet<Action>(actions.values());
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public TransitionModel transitionModelFor(State state, Action action) {
		TransitionModelKey key = new TransitionModelKey(state, action);
		if (this.transitions.containsKey(key)) {
			return transitions.get(key);
		} else {
			throw new IllegalArgumentException(String.format("Ação %s não disponível no estado %s.", action, state));
		}
	}

	public TransitionModel transitionModelFor(Object stateToken, Object actionToken) {
		return this.transitionModelFor(this.findState(stateToken), this.findAction(actionToken));
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public int getMaxIteractions() {
		return maxIteractions;
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public RewardModel rewardModel() {
		return rewardModel;
	}

	public DeclarableWorldModel withStates(String... tokens) {
		for (Object token : tokens) {
			this.withState(token);
		}
		return this;
	}

	public DeclarableWorldModel withStates(Object... tokens) {
		for (Object token : tokens) {
			this.withState(token);
		}
		return this;
	}

	public DeclarableWorldModel withState(Object token) {
		State state = new State(token);
		this.states.put(token, state);
		this.availableActionsAt.put(state, new HashSet<Action>());
		return this;
	}

	public DeclarableWorldModel withActions(Object... tokens) {
		for (Object token : tokens) {
			this.withAction(token);
		}
		return this;
	}

	public DeclarableWorldModel withAction(Object token) {
		this.actions.put(token, new Action(token));
		return this;
	}

	public TransitionModelBuilder withTransition() {
		return new TransitionModelBuilder();
	}

	public DeclarableWorldModel withDiscountFactor(Double discountFactor) {
		rewardModel.setDiscountFactor(discountFactor);
		return this;
	}

	public DeclarableWorldModel withDefaultReward(Double defaultReward) {
		rewardModel.setDefaultReward(defaultReward);
		return this;
	}

	public DeclarableWorldModel withDefaultCost(Double defaultCost) {
		rewardModel.setDefaultCost(defaultCost);
		return this;
	}

	public DeclarableWorldModel withReward(Object stateToken, double reward) {
		rewardModel.setReward(findState(stateToken), reward);
		return this;
	}

	public DeclarableWorldModel withCost(Object actionToken, double cost) {
		rewardModel.setCost(findAction(actionToken), cost);
		return this;
	}

	public DeclarableWorldModel withInitialState(Object stateToken) {
		this.initialState = this.findState(stateToken);
		return this;
	}

	public DeclarableWorldModel withMaxIteractions(int maxIteractions) {
		this.maxIteractions = maxIteractions;
		return this;
	}

	public DeclarableWorldModel withGoals(Object... stateTokens) {
		for (Object stateToken : stateTokens) {
			this.withGoal(stateToken);
		}
		return this;
	}

	public DeclarableWorldModel withGoal(Object stateToken) {
		this.goals.add(this.findState(stateToken));
		return this;
	}

	public State findState(Object token) {
		if (states.containsKey(token)) {
			return this.states.get(token);
		} else {
			throw new IllegalArgumentException("Estado não declarado: " + token);
		}
	}

	public Action findAction(Object token) {
		if (actions.containsKey(token)) {
			return this.actions.get(token);
		} else {
			throw new IllegalArgumentException("Ação não declarada: " + token);
		}
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public State initialState() {
		return initialState;
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public boolean isGoal(State state) {
		return this.goals.contains(state);
	}

	/**
	 * {@inheritDoc}
	 */
	@Override
	public Set<Action> availableActionsAt(State state) {
		if (this.availableActionsAt.containsKey(state)) {
			return Collections.unmodifiableSet(this.availableActionsAt.get(state));
		} else {
			return Collections.emptySet();
		}
	}

	private static class TransitionModelKey {
		private State state;
		private Action action;

		public TransitionModelKey(State state, Action action) {
			this.state = state;
			this.action = action;
		}

		@Override
		public int hashCode() {
			final int prime = 31;
			int result = 1;
			result = prime * result + ((state == null) ? 0 : state.hashCode());
			result = prime * result + ((action == null) ? 0 : action.hashCode());
			return result;
		}

		@Override
		public boolean equals(Object obj) {
			if (this == obj) {
				return true;
			}
			if (obj == null) {
				return false;
			}
			if (getClass() != obj.getClass()) {
				return false;
			}
			TransitionModelKey other = (TransitionModelKey) obj;
			if (state == null) {
				if (other.state != null) {
					return false;
				}
			} else if (!state.equals(other.state)) {
				return false;
			}
			if (action == null) {
				if (other.action != null) {
					return false;
				}
			} else if (!action.equals(other.action)) {
				return false;
			}
			return true;
		}
	}

	public class TransitionModelBuilder {

		private State state;

		private Action action;

		private TransitionModel.TransitionModelBuilder builder = TransitionModel.aTransitionModel();

		public TransitionModelBuilder on(Object stateToken) {
			this.state = findState(stateToken);
			return this;
		}

		public TransitionModelBuilder given(Object actionToken) {
			this.action = findAction(actionToken);
			return this;
		}

		public TransitionModelBuilder to(Object stateToken, Double probability) {
			builder.with(findState(stateToken), probability);
			return this;
		}

		public void build() {
			transitions.put(new TransitionModelKey(state, action), builder.build());
			availableActionsAt.get(state).add(action);
		}
	}
}
